import sensor, image, time, lcd, gc, cmath, struct
from maix import KPU #K210的AI加速器模块，用于深度学习
from modules import ybkey
from machine import UART
from fpioa_manager import fm

#初始化串口函数
def initUart():
    # binding UART2 IO:6->RX, 8->TX，引脚6为接收端，8为发送端
    fm.register(6, fm.fpioa.UART2_RX)
    fm.register(8, fm.fpioa.UART2_TX)
    yb_uart = UART(UART.UART2, 115200, 8, 0, 0, timeout=1000, read_buf_len=5)
    return yb_uart

#串口数据读取函数，接收返回的金额信息并显示
def readUart(yb_uart):
    while True:
        if yb_uart.any():#检查串口接收缓冲区中是否有数据
            read_data = yb_uart.read()#读取数据
            if read_data == None:
                continue
            return read_data
#写串口函数
def writeUart(yb_uart, payload):
    # 数据长度（前 4 个字节，大端序）
    length = len(payload)  #数据长度
    length_bytes = struct.pack('>I', length)  # '>I' 表示大端无符号整数 (4 字节)
    # 封装数据：长度（占4字节） + 有效数据（知道数据长度后再读能够确保所读数据的完整性）
    data_to_send = length_bytes + payload
    yb_uart.write(data_to_send)#封装好的数据发送到串口

#初始化屏幕，开机为红色
lcd.init()                          # Init lcd display
lcd.clear(lcd.RED)                  # Clear lcd screen.

#设置摄像头：复位，设置像素格式和分辨率并等待摄像头稳定
sensor.reset()                      # Reset and initialize the sensor.
sensor.set_pixformat(sensor.RGB565) # Set pixel format to RGB565 (or GRAYSCALE)
sensor.set_framesize(sensor.QVGA)   # Set frame size to QVGA (320x240)

sensor.skip_frames(time = 1000)     # Wait for settings take effect.

clock = time.clock()                # 1.Create a clock object to track the FPS.：fps是衡量系统处理图像的速度，显示系统当前的实时帧率，是判断系统性能的重要指标：<10慢，>30流畅

print("ready load model")

labels = ["zheng", "shousi", "yerong", "zimi", "rousong"] #类名称，按照label.txt顺序填写
anchor = (2.19, 1.94, 2.38, 2.53, 3.00, 2.38, 3.09, 3.25, 4.56, 4.31) #训练得到的anchor.txt中的五个锚点，预测五个面包边界框的(宽，高)
isPress = False  #按键初始未被按下
kpu = KPU()#创建kpu对象
# 从sd卡中加载模型
kpu.load_kmodel('/sd/det.kmodel')
#模型初始化
kpu.init_yolo2(
anchor, #预定义锚点列表
anchor_num=5, #使用的锚点数量（5个）
img_w=320, img_h=240, #图像宽高就是摄像头分辨率，看到多大就多大
net_w=320 , net_h=240 ,
layer_w=10 ,layer_h=8, #图像划分的网格大小，每个网格负责预测目标中心落在该区域内的对象
threshold=0.6, #保留概率大于0.6的预测框
nms_value=0.3, #合并重叠度超过30%的预测框
classes=len(labels)#待检测的目标类别数与所给的labels保持一致
)

key = ybkey() #初始化按键对象
yb_uart = initUart() #调用函数，初始化串口
while(True):
    gc.collect()  #回收内存

    clock.tick()  #2.开始计时（计算FPS）
    img = sensor.snapshot() #摄像头拍摄一张图像
    keyState = key.is_press(); #检测按键是否按下
    if keyState == 0: #未按下则继续显示图像并持续循环
        lcd.display(img)
        continue

    #如果检测到按键按下，并防抖处理
    print("have key press")
    time.sleep(0.5)
    isPress = True
    # 因为按键按下，所以将拍摄到的图像输入到kpu进行模型推理
    kpu.run_with_output(img)
    #获取检测结果：dect的数据结构：每个元素是一个元组，格式为（x,y,w,h,class_id,confidence）
    dect = kpu.regionlayer_yolo2()

    fps = clock.fps() #3.计算当前的FPS

    #如果检测到物体，dect>0就是有框
    if len(dect) > 0:
        breads = [0, 0, 0, 0, 0] #最多五个，先预设一下五个大的数组

        # for循环遍历每个检测结果并在图像上画框并显示标签及置信度
        for l in dect :  #遍历dect的每个检测结果

           # 统计检测出的每个类别面包的数量
            breads[l[4]] += 1  #l[4]为class_id,对应该类别在labels数组中的下标

            # 在图像上 用前四个元素绘制框，即（x,y,w,h）
            a = img.draw_rectangle(l[0],l[1],l[2],l[3],color=(0,255,0))
            # 在框上左上角（坐标（x,y））显示标签和置信度
            info = "%s %.3f" % (labels[l[4]], l[5])
            a = img.draw_string(l[0],l[1],info,color=(255,0,0),scale=2.0)

            print(info)
            del info #每次循环清空info,防止数据累积

        # 4.显示fps
        a = img.draw_string(0, 0, "%2.1ffps" %(fps),color=(0,60,255),scale=2.0)
        info = ""
        # 循环遍历每个检测结果并拼接成一整个字符串，到info上
        for i, bread in enumerate(breads):  #i是数组下标，bread是breads数组的i下标对应的值：即面包个数
            if bread > 0: # 面包个数大于0则执行拼接
                info += labels[i] + ":" + str(bread) + ';'
        print(info)
        writeUart(yb_uart, info);  #把检测到的面包数据即拼接好的数据Info发送到串口

        lcd.display(img) #最后lcd上显示带检测结果的图像

        money = readUart(yb_uart)  #java那边处理后会给串口发一个金额数据，这边阻塞等待读串口
        money = "total cost: " + str(money)  # 读到金额，输出到lcd上
        lcd.clear(lcd.BLUE)  # 清屏成蓝色，后显示白色金额数据money
        lcd.draw_string(100, 100, money, lcd.WHITE, lcd.BLUE)
        print(money)
    else:  #如果未检测到面包，则显示蓝屏并显示检测失败
        lcd.clear(lcd.BLUE)
        lcd.draw_string(100, 100, "detect error", lcd.WHITE, lcd.BLUE)
    # 等待按键再次按下以继续外层的while循环
    while True:
        flag = key.is_press()
        if flag == 1:
            print("have key press2")
            time.sleep(0.5)
            break
